import streamlit as st
import yaml
import os

def save_to_yaml(data, filename='config.yaml'):
    with open(filename, 'w') as file:
        yaml.dump(data, file)

def main():
    st.title("Model Parameters Configuration")

    # 验证相关参数
    st.header("Validation Parameters")
    batch_size_val = st.number_input("Batch Size (Validation)", min_value=1, value=32, step=1)
    num_workers_val = st.number_input("Number of Workers (Validation)", min_value=0, value=4, step=1)
    val_split = st.slider("Validation Split", min_value=0.0, max_value=1.0, value=0.2, step=0.05)

    # 推理相关参数
    st.header("Inference Parameters")
    batch_size_inf = st.number_input("Batch Size (Inference)", min_value=1, value=32, step=1)
    num_workers_inf = st.number_input("Number of Workers (Inference)", min_value=0, value=4, step=1)
    threshold = st.slider("Threshold", min_value=0.0, max_value=1.0, value=0.5, step=0.01)

    # 保存配置按钮
    if st.button("Save Configuration"):
        config = {
            "validation": {
                "batch_size": batch_size_val,
                "num_workers": num_workers_val,
                "val_split": val_split
            },
            "inference": {
                "batch_size": batch_size_inf,
                "num_workers": num_workers_inf,
                "threshold": threshold
            }
        }
        save_to_yaml(config)
        st.success("Configuration saved to config.yaml")

    # 显示配置参数
    if st.button("Show Configuration"):
        st.write("### Current Configuration")
        st.write(f"**Validation Batch Size:** {batch_size_val}")
        st.write(f"**Validation Number of Workers:** {num_workers_val}")
        st.write(f"**Validation Split:** {val_split}")
        st.write(f"**Inference Batch Size:** {batch_size_inf}")
        st.write(f"**Inference Number of Workers:** {num_workers_inf}")
        st.write(f"**Threshold:** {threshold}")

if __name__ == "__main__":
    main()